import numpy as np

class BallTreeNode:
    """
    定义BallTree的节点结构
    """
    def __init__(self, points, indices):
        self.points = points    # 当前节点包含的数据点
        self.indices = indices    # 数据点的索引
        self.center = np.mean(points, axis=0)    # 节点中心
        self.radius = np.max(np.linalg.norm(points - self.center, axis=1))  # 包围球半径
        self.left = None    # 左子节点
        self.right = None    # 右子节点

class BallTree:
    """
    BallTree实现
    """
    def __init__(self, data, leaf_size=2):
        self.data = data
        self.leaf_size = leaf_size  # 叶子节点的最小大小
        self.root = self._build_tree(data, np.arange(len(data)))
    
    def _build_tree(self, points, indices):
        if len(points) <= self.leaf_size:
            return BallTreeNode(points, indices)
        
        p1, p2 = self._find_farthest_points(points)
        
        left_indices = []
        right_indices = []
        
        for idx, point in enumerate(points):
            if np.linalg.norm(point - points[p1]) < np.linalg.norm(point - points[p2]):
                left_indices.append(idx)
            else:
                right_indices.append(idx)
        
        left_points = points[left_indices]
        right_points = points[right_indices]
        
        node = BallTreeNode(points, indices)
        node.left = self._build_tree(left_points, indices[left_indices])
        node.right = self._build_tree(right_points, indices[right_indices])
        return node

    def _find_farthest_points(self, points):
        max_distance = -1
        p1, p2 = 0, 0
        
        for i in range(len(points)):
            for j in range(i + 1, len(points)):
                distance = np.linalg.norm(points[i] - points[j])
                if distance > max_distance:
                    max_distance = distance
                    p1, p2 = i, j
        return p1, p2

    def query(self, query_point, k=1):
        """
        查询最接近的k个点
        """
        self.k = k
        self.best = [] # 保存当前最接近的k个点
        self._search(self.root, query_point)
        return sorted(self.best, key=lambda x: x[0]) # 按距离排序

    def _search(self, node, query_point):
        if node is None:
            return
        
        # 计算查询到当前节点中心的距离
        center_distance = np.linalg.norm(query_point - node.center)
        if len(self.best) >= self.k and center_distance - node.radius > self.best[-1][0]:
            return # 如果当前节点不可能包含更近的点，直接返回
        
        if node.left is None and node.right is None: # 叶子节点
            for i, point in enumerate(node.points):
                distance = np.linalg.norm(query_point - point)
                if len(self.best) < self.k:
                    self.best.append((distance, node.indices[i]))
                else:
                    if distance < self.best[-1][0]:
                        self.best[-1] = (distance, node.indices[i])
                self.best = sorted(self.best, key=lambda x: x[0]) # 保持按距离排序
            return
        
        # 递归搜索左右子节点
        self._search(node.left, query_point)
        self._search(node.right, query_point)

# 数据生成
np.random.seed(42)
data = np.random.rand(10, 3) # 生成10个三维数据点

# 构建BallTree
tree = BallTree(data)

# 查询最近邻
query_point = np.array([0.5, 0.5, 0.5]) # 查询点
k = 3 # 查询最近的3个点
results = tree.query(query_point, k=k)

# 输出结果
print(f"查询点: {query_point}")
print("最近邻结果:")
for distance, index in results:
    print(f"索引: {index}, 距离: {distance:.6f}, 坐标: {data[index]}")